from pprint import pprint

import numpy as np


# n 因子数，level 水平数， key 重点关注因素
def latin_square(n=3, level=5, key=0):
    assert n == 3, "因素水量目前只支持3"
    assert level >= 3, "水平数不能小于3"

    keys = [0, 1]
    for i in range(2, level):
        if i & 1:
            keys.append(int(i / 2 + 1))
        else:
            keys.append(level - int(i / 2))
    if level & 1:
        keys.reverse()

    res = []
    for i in range(level):
        for j in range(level):
            index = []
            if key == 0:
                index = [keys[j], i, j]
            elif key == 1:
                index = [i, keys[j], j]
            elif key == 2:
                index = [i, j, keys[j]]
            res.append(index)
        for j in range(len(keys)):
            keys[j] = (keys[j] + 1) % level

    pprint(len(res))
    pprint(res)
    return res


if __name__ == '__main__':
    res = latin_square(level=6)

